import numpy as np
from collections import defaultdict

import torch
from torch import nn
import torch.nn.functional as F
from util.box_ops import box_cxcywh_to_xyxy


class VCOCOEvaluator():

    def __init__(self, preds, gts, correct_mat, args):
        self.max_hois = args.max_hois
        self.fp = defaultdict(list)       
        self.tp = defaultdict(list)
        self.score = defaultdict(list)
        self.sum_gts = defaultdict(lambda: 0)

        self.use_nms_filter = args.use_nms_filter
        self.nms_thresh = args.nms_thresh
        self.nms_alpha = args.nms_alpha
        self.nms_beta = args.nms_beta
        self.HO_scale = args.HO_scale
        self.query_thres = args.query_thres


        self.obj_fp = defaultdict(list)       
        self.obj_tp = defaultdict(list)
        self.obj_score = defaultdict(list)
        self.obj_sum_gts = defaultdict(lambda: 0)


        self.fp_hm_obj_bbox = defaultdict(list)
        self.fp_hm_bbox = defaultdict(list)
        self.fp_obj_bbox = defaultdict(list)
        self.fp_obj_verb = defaultdict(list)
        self.fp_hm_verb = defaultdict(list)
        self.fp_verb = defaultdict(list)
        self.fp_hm_obj_verb = defaultdict(list)
        self.fp_repeat = defaultdict(list)
        self.fp_other = defaultdict(list)

        self.verb_classes = ['hold_obj', 'stand', 'sit_instr', 'ride_instr', 'walk', 'look_obj', 'hit_instr', 'hit_obj',
                             'eat_obj', 'eat_instr', 'jump_instr', 'lay_instr', 'talk_on_phone_instr', 'carry_obj',
                             'throw_obj', 'catch_obj', 'cut_instr', 'cut_obj', 'run', 'work_on_computer_instr',
                             'ski_instr', 'surf_instr', 'skateboard_instr', 'smile', 'drink_instr', 'kick_obj',
                             'point_instr', 'read_obj', 'snowboard_instr']
        #   stand walk run smile point_instr
        self.thesis_map_indices = [0, 2, 3, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 20, 21, 22, 24, 25, 27, 28]

        self.preds = []
        for img_preds in preds:
            img_preds = {k: v.to('cpu').numpy() for k, v in img_preds.items()}
            bboxes = [{'bbox': bbox, 
                       'category_id': label, 
                       'score':score} for bbox, label, score \
                in zip(img_preds['boxes'], img_preds['labels'], img_preds['score'])]
            #   100, 29
            hoi_scores = img_preds['verb_scores'] #* (img_preds['obj_score'] ** self.HO_scale) * (img_preds['sub_score'] ** self.HO_scale)
            verb_labels = np.tile(np.arange(hoi_scores.shape[1]), (hoi_scores.shape[0], 1))
            subject_ids = np.tile(img_preds['sub_ids'], (hoi_scores.shape[1], 1)).T
            object_ids = np.tile(img_preds['obj_ids'], (hoi_scores.shape[1], 1)).T
            #   2900, 
            hoi_scores = hoi_scores.ravel()
            verb_labels = verb_labels.ravel()
            subject_ids = subject_ids.ravel()
            object_ids = object_ids.ravel()

            if len(subject_ids) > 0:
                object_labels = np.array([bboxes[object_id]['category_id'] for object_id in object_ids])
                #   语义级别过滤矩阵
                correct_mat = np.concatenate((correct_mat, np.ones((correct_mat.shape[0], 1))), axis=1)
                masks = correct_mat[verb_labels, object_labels]
                hoi_scores *= masks
                hoi_scores[hoi_scores < self.query_thres] = 0.
                new_hoi_scores = hoi_scores[hoi_scores>0.] 
                verb_labels = verb_labels[hoi_scores>0.]
                object_ids =  object_ids[hoi_scores>0.]
                subject_ids = subject_ids[hoi_scores>0.]
                
                hois = [{'subject_id': subject_id, 
                         'object_id': object_id, 
                         'category_id': category_id, 
                         'score': score} for
                        subject_id, object_id, category_id, score in zip(subject_ids, object_ids, verb_labels, new_hoi_scores)]
                # hois.sort(key=lambda k: (k.get('score', 0)), reverse=True)
                # hois = hois[:self.max_hois]
            else:
                hois = []
            current_result = {'predictions': bboxes, 
                              'hoi_prediction': hois}    
            if self.use_nms_filter:
                current_result = self.triplet_nms_filter(current_result)
            current_result['hoi_prediction'].sort(key=lambda k: (k.get('score', 0)), reverse=True)
            current_result['hoi_prediction'] = current_result['hoi_prediction'][:self.max_hois]   
            self.preds.append(current_result)



                 
        self.gts = []
        for img_gts in gts:
            img_gts = {k: v.to('cpu').numpy() for k, v in img_gts.items() if k != 'id' and k != 'img_id' and k != 'filename'}
            self.gts.append({
                'annotations': [{'bbox': bbox, 
                                'category_id': label} for bbox, label in zip(img_gts['boxes'], img_gts['labels'])],
                'hoi_annotation': [{'subject_id': hoi[0], 
                                    'object_id': hoi[1], 
                                    'category_id': hoi[2]} for hoi in img_gts['hois']]
            })
            for hoi in self.gts[-1]['hoi_annotation']:
                self.sum_gts[hoi['category_id']] += 1
            for obj in self.gts[-1]['annotations']:
                self.obj_sum_gts[obj['category_id']] += 1

    def evaluate(self):
        for img_preds, img_gts in zip(self.preds, self.gts):#   every img
            pred_bboxes = img_preds['predictions']
            gt_bboxes = img_gts['annotations']
            
            pred_hois = img_preds['hoi_prediction']
            gt_hois = img_gts['hoi_annotation']
            
            if len(gt_bboxes) != 0:
                bbox_pairs, bbox_overlaps = self.compute_iou_mat(gt_bboxes, pred_bboxes)
                self.compute_fptp(pred_hois, gt_hois, bbox_pairs, pred_bboxes, bbox_overlaps)
            else:#  
                for pred_hoi in pred_hois:  #   every anno(100)
                    if isinstance(pred_hoi['category_id'], str):
                        pred_hoi['category_id'] = int(pred_hoi['category_id'].replace('\n', ''))
                    self.tp[pred_hoi['category_id']].append(0)
                    self.fp[pred_hoi['category_id']].append(1)
                    self.score[pred_hoi['category_id']].append(pred_hoi['score'])
                for pred_bbox in pred_bboxes:   #   every hm, obj(200)
                    if isinstance(pred_bbox['category_id'], str):
                        pred_bbox['category_id'] = int(pred_bbox['category_id'].replace('\n', ''))
                    self.obj_tp[pred_bbox['category_id']].append(0)
                    self.obj_fp[pred_bbox['category_id']].append(1)
                    self.obj_score[pred_bbox['category_id']].append(pred_bbox['score'])
        
        map = self.compute_map()
        return map


    def compute_iou_mat(self, bbox_list1, bbox_list2, overlap_iou=0.5):
        iou_mat = np.zeros((len(bbox_list1), len(bbox_list2)))
        if len(bbox_list1) == 0 or len(bbox_list2) == 0:
            return {}
        for i, bbox1 in enumerate(bbox_list1):
            for j, bbox2 in enumerate(bbox_list2):
                iou_i = self.compute_IOU(bbox1, bbox2)
                iou_mat[i, j] = iou_i

        iou_mat_ov=iou_mat.copy()
        iou_mat[iou_mat >= overlap_iou] = 1
        iou_mat[iou_mat < overlap_iou] = 0

        match_pairs = np.nonzero(iou_mat)# value , index
        match_pairs_dict = {}
        match_pair_overlaps = {}
        if iou_mat.max() > 0:
            for i, pred_id in enumerate(match_pairs[1]):#   match
                if pred_id not in match_pairs_dict.keys():
                    match_pairs_dict[pred_id] = []
                    match_pair_overlaps[pred_id]=[]
                match_pairs_dict[pred_id].append(match_pairs[0][i])
                match_pair_overlaps[pred_id].append(iou_mat_ov[match_pairs[0][i],pred_id])
        ######
        pos_pred_ids = match_pairs_dict.keys()
        #   一对多，只能计算precision
        #bbox_list2.sort(key=lambda k: (k.get('score', 0)), reverse=True)
        for pred_id, pred in enumerate(bbox_list2):#0-200
            if isinstance(pred['category_id'], str):
                pred['category_id'] = int(pred['category_id'].replace('\n', ''))
            # if pred_id in pos_pred_ids:
            #     self.obj_fp[pred['category_id']].append(0)
            #     self.obj_tp[pred['category_id']].append(1)      
            # else:
            #     self.obj_fp[pred['category_id']].append(1)
            #     self.obj_tp[pred['category_id']].append(0)
            # self.obj_score[pred['category_id']].append(pred['score'])              
        return match_pairs_dict, match_pair_overlaps

    def compute_IOU(self, bbox1, bbox2):
        if isinstance(bbox1['category_id'], str):
            bbox1['category_id'] = int(bbox1['category_id'].replace('\n', ''))
        if isinstance(bbox2['category_id'], str):
            bbox2['category_id'] = int(bbox2['category_id'].replace('\n', ''))
        if bbox1['category_id'] == bbox2['category_id']:
            rec1 = bbox1['bbox']
            rec2 = bbox2['bbox']
            # computing area of each rectangles
            S_rec1 = (rec1[2] - rec1[0]+1) * (rec1[3] - rec1[1]+1)
            S_rec2 = (rec2[2] - rec2[0]+1) * (rec2[3] - rec2[1]+1)

            # computing the sum_area
            sum_area = S_rec1 + S_rec2

            # find the each edge of intersect rectangle
            left_line = max(rec1[1], rec2[1])
            right_line = min(rec1[3], rec2[3])
            top_line = max(rec1[0], rec2[0])
            bottom_line = min(rec1[2], rec2[2])
            # judge if there is an intersect
            if left_line >= right_line or top_line >= bottom_line:
                return 0
            else:
                intersect = (right_line - left_line+1) * (bottom_line - top_line+1)
                return intersect / (sum_area - intersect)
        else:
            return 0


    def compute_fptp(self, pred_hois, gt_hois, match_pairs, pred_bboxes, bbox_overlaps):
        pos_pred_ids = match_pairs.keys()
        vis_tag = np.zeros(len(gt_hois))
        pred_hois.sort(key=lambda k: (k.get('score', 0)), reverse=True)
      
        if len(pred_hois) != 0:
            for pred_hoi in pred_hois:
                # if pred_hoi['score']<=0.:
                #     continue
                is_obj_match_sub = 0
                is_obj_match_obj = 0

                is_match = 0
                fp_hm_obj_bbox = 0
                fp_hm_bbox = 0
                fp_obj_bbox = 0
                fp_obj_verb = 0
                fp_hm_verb = 0
                fp_verb = 0
                fp_hm_obj_verb = 0
                #####
                
                max_overlap = 0
                max_gt_hoi = 0
                for gt_hoi in gt_hois:
                    if len(match_pairs) != 0 and \
                    pred_hoi['subject_id'] in pos_pred_ids and \
                    gt_hoi['object_id'] == -1:          #   scenario_2
                        pred_sub_ids = match_pairs[pred_hoi['subject_id']]
                        pred_sub_overlaps = bbox_overlaps[pred_hoi['subject_id']]
                        pred_category_id = pred_hoi['category_id']
                        if gt_hoi['subject_id'] in pred_sub_ids and \
                        pred_category_id == gt_hoi['category_id']:
                            is_match = 1
                            is_obj_match_sub = 1
                            
                            min_overlap_gt = pred_sub_overlaps[pred_sub_ids.index(gt_hoi['subject_id'])]
                            if min_overlap_gt > max_overlap:
                                max_overlap = min_overlap_gt
                                max_gt_hoi = gt_hoi
                        elif gt_hoi['subject_id'] in pred_sub_ids:
                            if pred_category_id != gt_hoi['category_id']:
                                fp_verb = 1 
                        else:
                            if pred_category_id == gt_hoi['category_id']:
                                fp_hm_bbox = 1
                            else:
                                fp_hm_verb = 1
                    
                    #   匹配到物体的
                    elif len(match_pairs) != 0 and \
                    pred_hoi['subject_id'] in pos_pred_ids and \
                    pred_hoi['object_id'] in pos_pred_ids:  #   scenario_1
                        pred_sub_ids = match_pairs[pred_hoi['subject_id']]
                        pred_obj_ids = match_pairs[pred_hoi['object_id']]
                        pred_sub_overlaps = bbox_overlaps[pred_hoi['subject_id']]
                        pred_obj_overlaps = bbox_overlaps[pred_hoi['object_id']]
                        pred_category_id = pred_hoi['category_id']
                        if gt_hoi['subject_id'] in pred_sub_ids and \
                        gt_hoi['object_id'] in pred_obj_ids and \
                        pred_category_id == gt_hoi['category_id']:  
                            is_match = 1
                            
                            is_obj_match_sub = 1
                            is_obj_match_obj = 1
                            
                            min_overlap_gt = min(pred_sub_overlaps[pred_sub_ids.index(gt_hoi['subject_id'])],
                                                 pred_obj_overlaps[pred_obj_ids.index(gt_hoi['object_id'])])    
                            if min_overlap_gt > max_overlap:
                                max_overlap = min_overlap_gt
                                max_gt_hoi = gt_hoi
                        
                        elif gt_hoi['subject_id'] in pred_sub_ids:
                            is_obj_match_sub = 1
                            
                            if gt_hoi['object_id'] in pred_obj_ids:
                                is_obj_match_obj = 1
                                
                                if pred_category_id != gt_hoi['category_id']:
                                    fp_verb = 1
                            else:
                                if pred_category_id != gt_hoi['category_id']:
                                    fp_obj_verb = 1
                                else:
                                    fp_obj_bbox = 1
                        else:
                            if gt_hoi['object_id'] in pred_obj_ids:
                                
                                is_obj_match_obj = 1
                                
                                if pred_category_id != gt_hoi['category_id']:
                                    fp_hm_verb = 1
                                else:
                                    fp_hm_bbox = 1
                            else: 
                                if pred_category_id != gt_hoi['category_id']: 
                                    fp_hm_obj_verb = 1
                                else: 
                                    fp_hm_obj_bbox = 1
                                              
                    #   没匹到人物
                    elif len(match_pairs) != 0 and \
                    pred_hoi['subject_id'] not in pos_pred_ids and \
                    pred_hoi['object_id'] in pos_pred_ids:
                        pred_obj_ids = match_pairs[pred_hoi['object_id']]
                        pred_category_id = pred_hoi['category_id']
                        if gt_hoi['object_id'] in pred_obj_ids:   
                            is_obj_match_obj = 1                         
                            if pred_category_id != gt_hoi['category_id']:
                                fp_hm_verb = 1
                            else: 
                                fp_hm_bbox = 1
                        else:   
                            if pred_category_id != gt_hoi['category_id']:
                                fp_hm_obj_verb = 1
                            else: 
                                fp_hm_obj_verb = 1 

                    #   没匹到物体
                    elif len(match_pairs) != 0 and \
                    pred_hoi['subject_id'] in pos_pred_ids and \
                    pred_hoi['object_id'] not in pos_pred_ids:  
                        pred_sub_ids = match_pairs[pred_hoi['subject_id']]
                        pred_category_id = pred_hoi['category_id']
                        if gt_hoi['subject_id'] in pred_sub_ids: 
                            is_obj_match_sub = 1
                            if pred_category_id != gt_hoi['category_id']:
                                fp_obj_verb = 1
                            else: 
                                fp_obj_bbox = 1
                        else: 
                            if pred_category_id != gt_hoi['category_id']:
                                fp_hm_obj_verb = 1
                            else: 
                                fp_hm_obj_bbox = 1 

                    #   没匹到人物和物体
                    elif len(match_pairs) != 0 and \
                    pred_hoi['subject_id'] not in pos_pred_ids and \
                    pred_hoi['object_id'] not in pos_pred_ids:  
                        pred_category_id = pred_hoi['category_id']
                        if pred_category_id != gt_hoi['category_id']:
                            fp_hm_obj_verb = 1
                        else:
                            fp_hm_obj_bbox = 1
                
                            
                if is_match == 1 and vis_tag[gt_hois.index(max_gt_hoi)] == 0:
                    self.fp[pred_hoi['category_id']].append(0)
                    self.tp[pred_hoi['category_id']].append(1)
                    vis_tag[gt_hois.index(max_gt_hoi)] = 1
                    
                    self.fp_hm_obj_bbox[pred_hoi['category_id']].append(0)
                    self.fp_hm_bbox[pred_hoi['category_id']].append(0)
                    self.fp_obj_bbox[pred_hoi['category_id']].append(0)
                    self.fp_obj_verb[pred_hoi['category_id']].append(0)
                    self.fp_hm_verb[pred_hoi['category_id']].append(0)
                    self.fp_verb[pred_hoi['category_id']].append(0)
                    self.fp_hm_obj_verb[pred_hoi['category_id']].append(0)
                    self.fp_repeat[pred_hoi['category_id']].append(0)
                    self.fp_other[pred_hoi['category_id']].append(0)


                    self.obj_fp[pred_bboxes[pred_hoi['subject_id']]['category_id']].append(0)
                    self.obj_tp[pred_bboxes[pred_hoi['subject_id']]['category_id']].append(1)
                    
                    self.obj_fp[pred_bboxes[pred_hoi['object_id']]['category_id']].append(0)
                    self.obj_tp[pred_bboxes[pred_hoi['object_id']]['category_id']].append(1)
                    
                else:    
                    
                    self.fp[pred_hoi['category_id']].append(1)
                    self.tp[pred_hoi['category_id']].append(0)

                    if fp_hm_obj_bbox == 1:
                        self.fp_hm_obj_bbox[pred_hoi['category_id']].append(1)
                        self.fp_hm_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_verb[pred_hoi['category_id']].append(0)
                        self.fp_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_repeat[pred_hoi['category_id']].append(0)
                        self.fp_other[pred_hoi['category_id']].append(0)              
                    elif fp_hm_bbox == 1:
                        self.fp_hm_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_hm_bbox[pred_hoi['category_id']].append(1)
                        self.fp_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_verb[pred_hoi['category_id']].append(0)
                        self.fp_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_repeat[pred_hoi['category_id']].append(0)  
                        self.fp_other[pred_hoi['category_id']].append(0)
                    elif fp_obj_bbox == 1:
                        self.fp_hm_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_hm_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_bbox[pred_hoi['category_id']].append(1)
                        self.fp_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_verb[pred_hoi['category_id']].append(0)
                        self.fp_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_obj_verb[pred_hoi['category_id']].append(0) 
                        self.fp_repeat[pred_hoi['category_id']].append(0)    
                        self.fp_other[pred_hoi['category_id']].append(0)                   
                    elif fp_obj_verb == 1:
                        self.fp_hm_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_hm_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_verb[pred_hoi['category_id']].append(1)
                        self.fp_hm_verb[pred_hoi['category_id']].append(0)
                        self.fp_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_obj_verb[pred_hoi['category_id']].append(0)  
                        self.fp_repeat[pred_hoi['category_id']].append(0)
                        self.fp_other[pred_hoi['category_id']].append(0)
                    elif fp_hm_verb == 1:
                        self.fp_hm_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_hm_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_verb[pred_hoi['category_id']].append(1)
                        self.fp_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_repeat[pred_hoi['category_id']].append(0)  
                        self.fp_other[pred_hoi['category_id']].append(0)
                    elif fp_verb == 1:
                        self.fp_hm_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_hm_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_verb[pred_hoi['category_id']].append(0)
                        self.fp_verb[pred_hoi['category_id']].append(1)
                        self.fp_hm_obj_verb[pred_hoi['category_id']].append(0) 
                        self.fp_repeat[pred_hoi['category_id']].append(0) 
                        self.fp_other[pred_hoi['category_id']].append(0)
                    elif fp_hm_obj_verb == 1:
                        self.fp_hm_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_hm_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_verb[pred_hoi['category_id']].append(0)
                        self.fp_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_obj_verb[pred_hoi['category_id']].append(1)  
                        self.fp_repeat[pred_hoi['category_id']].append(0)
                        self.fp_other[pred_hoi['category_id']].append(0)
                    elif is_match == 1:
                        self.fp_hm_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_hm_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_verb[pred_hoi['category_id']].append(0)
                        self.fp_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_obj_verb[pred_hoi['category_id']].append(0)  
                        self.fp_repeat[pred_hoi['category_id']].append(1)  
                        self.fp_other[pred_hoi['category_id']].append(0) 
                    else:                     
                        self.fp_hm_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_hm_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_bbox[pred_hoi['category_id']].append(0)
                        self.fp_obj_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_verb[pred_hoi['category_id']].append(0)
                        self.fp_verb[pred_hoi['category_id']].append(0)
                        self.fp_hm_obj_verb[pred_hoi['category_id']].append(0)  
                        self.fp_repeat[pred_hoi['category_id']].append(0)  
                        self.fp_other[pred_hoi['category_id']].append(1)    
                                
                    if is_obj_match_sub == 1:
                        self.obj_fp[pred_bboxes[pred_hoi['subject_id']]['category_id']].append(0)
                        self.obj_tp[pred_bboxes[pred_hoi['subject_id']]['category_id']].append(1)   
                    else:   
                        self.obj_fp[pred_bboxes[pred_hoi['subject_id']]['category_id']].append(1)
                        self.obj_tp[pred_bboxes[pred_hoi['subject_id']]['category_id']].append(0)  
                    if is_obj_match_obj == 1:
                        self.obj_fp[pred_bboxes[pred_hoi['object_id']]['category_id']].append(0)
                        self.obj_tp[pred_bboxes[pred_hoi['object_id']]['category_id']].append(1)   
                    else:   
                        self.obj_fp[pred_bboxes[pred_hoi['object_id']]['category_id']].append(1)
                        self.obj_tp[pred_bboxes[pred_hoi['object_id']]['category_id']].append(0)

                self.score[pred_hoi['category_id']].append(pred_hoi['score'])
                self.obj_score[pred_bboxes[pred_hoi['subject_id']]['category_id']].append(pred_bboxes[pred_hoi['subject_id']]['score'])          
                self.obj_score[pred_bboxes[pred_hoi['object_id']]['category_id']].append(pred_bboxes[pred_hoi['object_id']]['score'])             
                




    def compute_map(self):
        print('------------------------------------------------------------')
        ap = defaultdict(lambda: 0)
        
        recall = defaultdict(lambda: 0)
        precision = defaultdict(lambda: 0)
        acc = defaultdict(lambda: 0)
        F1_score = defaultdict(lambda: 0)

        obj_ap = defaultdict(lambda: 0)
        obj_recall = defaultdict(lambda: 0)
        obj_precision = defaultdict(lambda: 0)
        obj_acc= defaultdict(lambda: 0)    
        obj_F1_score = defaultdict(lambda: 0)

        rate_fp_hm_obj_bbox = defaultdict(lambda: 0)
        rate_fp_hm_bbox = defaultdict(lambda: 0)
        rate_fp_obj_bbox = defaultdict(lambda: 0)  
        rate_fp_obj_verb = defaultdict(lambda: 0)
        rate_fp_hm_verb = defaultdict(lambda: 0)
        rate_fp_verb = defaultdict(lambda: 0)
        rate_fp_hm_obj_verb = defaultdict(lambda: 0)
        rate_fp_repeat = defaultdict(lambda: 0)
        rate_fp_other = defaultdict(lambda: 0)

        for category_id in sorted(list(self.obj_sum_gts.keys())):
            obj_sum_gts = self.obj_sum_gts[category_id]
            if obj_sum_gts == 0:
                continue
            obj_tp = np.array((self.obj_tp[category_id]))
            obj_fp = np.array((self.obj_fp[category_id]))
            
            tp_sum =  np.sum(obj_tp)
            fp_sum =  np.sum(obj_fp)  
            obj_precision[category_id] = tp_sum/(fp_sum + tp_sum) if (fp_sum + tp_sum)!=0 else 0 
            # if len(obj_tp) == 0:
            #     obj_ap[category_id] = 0
            # else:
            #     #   29
            #     score = np.array(self.obj_score[category_id])
            #     sort_inds = np.argsort(-score)
            #     tp = obj_tp[sort_inds]
            #     fp = obj_fp[sort_inds]
            #     #   (44006,)
            #     tp = np.cumsum(tp)
            #     fp = np.cumsum(fp)

            #     #   (44006,)
            #     rec = tp / obj_sum_gts               
            #     prec = tp / (fp + tp)
            #     obj_ap[category_id] = self.voc_ap(rec, prec)

            #     tp_sum = tp[-1]
            #     fp_sum =  fp[-1]           
            #     obj_recall[category_id] = tp_sum / obj_sum_gts 
            #     obj_precision[category_id] = tp_sum / (fp_sum + tp_sum)
            #     obj_acc[category_id] = tp_sum / (obj_sum_gts + fp_sum)      
            #     obj_F1_score[category_id] = (2*obj_recall[category_id] * obj_precision[category_id]) / (obj_recall[category_id] + obj_precision[category_id]) \
            #         if (obj_recall[category_id] + obj_precision[category_id])!=0 else 0
    
        # obj_ap_list = list(obj_ap.values())
        # m_obj_ap_all = np.mean(obj_ap_list)
        # m_obj_ap_human = np.mean(obj_ap_list)
        # m_obj_ap_object = np.mean(obj_ap_list)

        # m_obj_recall_list = list(obj_recall.values())
        # m_obj_recall_all = np.mean(m_obj_recall_list)
        # m_obj_recall_human = np.mean(m_obj_recall_list[0:1])
        # m_obj_recall_object = np.mean(m_obj_recall_list[1:])
               
        obj_precision_all_list = list(obj_precision.values())
        m_obj_precision_all = np.mean(obj_precision_all_list)
        m_obj_precision_human = np.mean(obj_precision_all_list[0:1])
        m_obj_precision_object = np.mean(obj_precision_all_list[1:])
        
        # m_obj_acc = np.mean(list(obj_acc.values()))
        # m_obj_F1_score = np.mean(list(obj_F1_score.values()))        
        


        aps = {}
        for category_id in sorted(list(self.sum_gts.keys())):
            #   3608
            sum_gts = self.sum_gts[category_id]
            if sum_gts == 0:
                continue
            #   (44006,)
            tp = np.array((self.tp[category_id]))
            fp = np.array((self.fp[category_id]))
        
            fp_hm_obj_bbox =  np.array((self.fp_hm_obj_bbox[category_id]))
            fp_hm_bbox = np.array((self.fp_hm_bbox[category_id]))
            fp_obj_bbox = np.array((self.fp_obj_bbox[category_id]))
            fp_obj_verb = np.array((self.fp_obj_verb[category_id]))
            fp_hm_verb = np.array((self.fp_hm_verb[category_id]))
            fp_verb = np.array((self.fp_verb[category_id]))
            fp_hm_obj_verb = np.array((self.fp_hm_obj_verb[category_id]))
            fp_repeat = np.array((self.fp_repeat[category_id]))
            fp_other = np.array((self.fp_other[category_id]))


            if len(tp) == 0:
                ap[category_id] = 0
            else:
                #   29
                score = np.array(self.score[category_id])
                sort_inds = np.argsort(-score)
                tp = tp[sort_inds]
                fp = fp[sort_inds]
                #   (44006,)
                tp = np.cumsum(tp)
                fp = np.cumsum(fp)

                #   (44006,)
                rec = tp / sum_gts               
                prec = tp / (fp + tp)
                ap[category_id] = self.voc_ap(rec, prec)

                tp_sum = tp[-1]
                fp_sum =  fp[-1]           
                recall[category_id] = tp_sum / sum_gts 
                precision[category_id] = tp_sum / (fp_sum + tp_sum)
                acc[category_id] = tp_sum / (sum_gts + fp_sum)      
                F1_score[category_id] = (2*recall[category_id] * precision[category_id]) / (recall[category_id] + precision[category_id]) \
                    if (recall[category_id] + precision[category_id])!=0 else 0
                    
                rate_fp_hm_obj_bbox[category_id] =  np.sum(fp_hm_obj_bbox)/(fp_sum + tp_sum)
                rate_fp_hm_bbox[category_id] =  np.sum(fp_hm_bbox)/(fp_sum + tp_sum)
                rate_fp_obj_bbox[category_id] =  np.sum(fp_obj_bbox)/(fp_sum + tp_sum)
                rate_fp_obj_verb[category_id] =  np.sum(fp_obj_verb)/(fp_sum + tp_sum)
                rate_fp_hm_verb[category_id] =  np.sum(fp_hm_verb)/(fp_sum + tp_sum)
                rate_fp_verb[category_id] =  np.sum(fp_verb)/(fp_sum + tp_sum)
                rate_fp_hm_obj_verb[category_id] =  np.sum(fp_hm_obj_verb)/(fp_sum + tp_sum)
                rate_fp_repeat[category_id] =  np.sum(fp_repeat)/(fp_sum + tp_sum)
                rate_fp_other[category_id] =  np.sum(fp_other)/(fp_sum + tp_sum)

            print('{:>23s}: #GTs = {:>04d}, AP = {:>.4f},\
recall = {:>.4f}, precision = {:>.4f},  acc = {:>.4f}, F1_score = {:>.4f},\
rate_fp_hm_obj_bbox = {:>.4f}, rate_fp_hm_bbox = {:>.4f}, rate_fp_obj_bbox = {:>.4f}, rate_fp_obj_verb = {:>.4f},\
rate_fp_hm_verb = {:>.4f}, rate_fp_verb = {:>.4f}, rate_fp_hm_obj_verb = {:>.4f}, rate_fp_repeat = {:>.4f}, rate_fp_other = {:>.4f}'.
                  format(self.verb_classes[category_id], sum_gts, ap[category_id], 
                         recall[category_id], precision[category_id], acc[category_id], F1_score[category_id],
                        rate_fp_hm_obj_bbox[category_id], rate_fp_hm_bbox[category_id], rate_fp_obj_bbox[category_id], rate_fp_obj_verb[category_id],
                        rate_fp_hm_verb[category_id], rate_fp_verb[category_id], rate_fp_hm_obj_verb[category_id], rate_fp_repeat[category_id], rate_fp_other[category_id]))
            aps['AP_{}'.format(self.verb_classes[category_id])] = ap[category_id]
        #   每一类的平均
        m_ap_all = np.mean(list(ap.values()))
        m_ap_thesis = np.mean([ap[category_id] for category_id in self.thesis_map_indices])
        
        m_recall = np.mean(list(recall.values()))
        m_precision = np.mean(list(precision.values()))
        m_acc = np.mean(list(acc.values()))
        m_F1_score = np.mean(list(F1_score.values()))
        
        m_fp_hm_obj_bbox =  np.mean(list(rate_fp_hm_obj_bbox.values()))
        m_fp_hm_bbox =  np.mean(list(rate_fp_hm_bbox.values()))
        m_fp_obj_bbox =  np.mean(list(rate_fp_obj_bbox.values()))
        m_fp_obj_verb =  np.mean(list(rate_fp_obj_verb.values()))
        m_fp_hm_verb =  np.mean(list(rate_fp_hm_verb.values()))
        m_fp_verb =  np.mean(list(rate_fp_verb.values()))
        m_fp_hm_obj_verb =  np.mean(list(rate_fp_repeat.values()))
        m_fp_repeat =  np.mean(list(rate_fp_repeat.values()))
        m_fp_other =  np.mean(list(rate_fp_other.values()))
        
        print('------------------------------------------------------------')
        print('mAP all: {:.4f} mAP thesis: {:.4f}\n\
m_recall: {:.4f}  m_precision: {:.4f}  m_acc: {:.4f}  m_F1_score: {:.4f},\n\
m_fp_hm_obj_bbox = {:>.4f} m_fp_hm_bbox = {:>.4f} m_fp_obj_bbox = {:>.4f} m_fp_obj_verb = {:>.4f} \
m_fp_hm_verb = {:>.4f} m_fp_verb = {:>.4f} m_fp_hm_obj_verb = {:>.4f} m_fp_repeat = {:>.4f} m_fp_other = {:>.4f}\n\
m_obj_precision_all= {:>.4f}, m_obj_precision_human= {:>.4f}, m_obj_precision_object= {:>.4f}\n\
'.format(
        m_ap_all, m_ap_thesis, 
        m_recall, m_precision, m_acc, m_F1_score,
        m_fp_hm_obj_bbox, m_fp_hm_bbox, m_fp_obj_bbox, m_fp_obj_verb, 
        m_fp_hm_verb, m_fp_verb, m_fp_hm_obj_verb, m_fp_repeat, m_fp_other,
        # m_obj_ap_all, m_obj_ap_human, m_obj_ap_object,
        # m_obj_recall_all, m_obj_recall_human, m_obj_recall_object,
        m_obj_precision_all, m_obj_precision_human, m_obj_precision_object))
        # m_obj_acc, m_obj_F1_score))
        print('------------------------------------------------------------')
        aps.update({'mAP_all': m_ap_all, 'mAP_thesis': m_ap_thesis, 
                    'm_recall':m_recall,'m_precision':m_precision, 'm_acc':m_acc,'m_F1_score':m_F1_score,
                    'm_fp_hm_obj_bbox':m_fp_hm_obj_bbox, "m_fp_hm_bbox":m_fp_hm_bbox, "m_fp_obj_bbox":m_fp_obj_bbox, "m_fp_obj_verb":m_fp_obj_verb,
                    "m_fp_hm_verb":m_fp_hm_verb, "m_fp_verb":m_fp_verb, "m_fp_hm_obj_verb":m_fp_hm_obj_verb, 'm_fp_repeat':m_fp_repeat, 'm_fp_other':m_fp_other,
                    # "m_obj_ap_all":m_obj_ap_all, "m_obj_ap_human":m_obj_ap_human, "m_obj_ap_object":m_obj_ap_object,
                    # "m_obj_recall_all":m_obj_recall_all, "m_obj_recall_human":m_obj_recall_human, "m_obj_recall_object":m_obj_recall_object,
                    "m_obj_precision_all":m_obj_precision_all, "m_obj_precision_human":m_obj_precision_human, "m_obj_precision_object":m_obj_precision_object,})
                    # "m_obj_acc":m_obj_acc, "m_obj_F1_score":m_obj_F1_score})
        return aps



    

    @torch.no_grad()
    def get_outputs(self, outputs, target_sizes):
        assert target_sizes.shape[1] == 2
        obj_scores, obj_labels = self.get_obj_cls(outputs, target_sizes)
        sub_boxes = self.get_sub_bbox(outputs, target_sizes)
        obj_boxes = self.get_obj_bbox(outputs, target_sizes)
        verb_scores = self.get_verb_score(outputs, target_sizes)
        matching_scores = self.get_matching_scores(outputs, target_sizes)
        return sub_boxes, obj_boxes, obj_scores, obj_labels, verb_scores, matching_scores
    

    def get_obj_cls(self,outputs, target_sizes):
        out_obj_logits = outputs['pred_obj_logits']
        assert len(out_obj_logits) == len(target_sizes)
        obj_prob = F.softmax(out_obj_logits, -1)
        obj_scores, obj_labels = obj_prob[..., :-1].max(-1)#change
        obj_scores[obj_labels==0]=0.
        return obj_scores, obj_labels

    def get_verb_score(self,outputs, target_sizes, topK=5):
        out_verb_logits = outputs['pred_verb_logits']
        verb_scores = out_verb_logits.sigmoid()
        #   100, 29
        _, verb_index  = torch.topk(-verb_scores, verb_scores.shape[-1]-topK, dim=-1)
        verb_scores = verb_scores.scatter(-1, verb_index, 0)
        return verb_scores

    def get_obj_bbox(self,outputs, target_sizes):
        out_obj_boxes = outputs['pred_obj_boxes']
        img_h, img_w = target_sizes.unbind(1)
        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(out_obj_boxes.device)
        obj_boxes = box_cxcywh_to_xyxy(out_obj_boxes)
        obj_boxes = obj_boxes * scale_fct[:, None, :]
        return obj_boxes

    def get_sub_bbox(self,outputs, target_sizes):
        out_sub_boxes = outputs['pred_sub_boxes']
        img_h, img_w = target_sizes.unbind(1)
        scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(out_sub_boxes.device)
        sub_boxes = box_cxcywh_to_xyxy(out_sub_boxes)
        sub_boxes = sub_boxes * scale_fct[:, None, :]  
        return sub_boxes

    def get_matching_scores(self,outputs, target_sizes):
        if self.use_matching:
            out_matching_logits = outputs['pred_matching_logits']
            matching_scores = out_matching_logits.sigmoid()[..., 0]
            #matching_scores = F.softmax(out_matching_logits, -1)[..., 0]
        else: 
            matching_scores = None
        return matching_scores



#############################################  PNMS ###########################################################
    @torch.no_grad()
    def triplet_nms_filter(self, preds):
        pred_od = preds['predictions']
        pred_hois = preds['hoi_prediction']
        all_triplets = {}
        for index, pred_hoi in enumerate(pred_hois):
            triplet = str(pred_od[pred_hoi['subject_id']]['category_id']) + '_' + \
                      str(pred_od[pred_hoi['object_id']]['category_id']) + '_' + \
                      str(pred_hoi['category_id'])

            if triplet not in all_triplets:
                all_triplets[triplet] = {'subs':[], 
                                         'objs':[], 
                                         'scores':[], 
                                         'indexes':[]}
            all_triplets[triplet]['subs'].append(pred_od[pred_hoi['subject_id']]['bbox'])
            all_triplets[triplet]['objs'].append(pred_od[pred_hoi['object_id']]['bbox'])
            all_triplets[triplet]['scores'].append(pred_hoi['score'])
            all_triplets[triplet]['indexes'].append(index)

        all_keep_inds = []
        for triplet, values in all_triplets.items():
            subs, objs, scores = values['subs'], values['objs'], values['scores']
            keep_inds = self.pairwise_nms(np.array(subs), np.array(objs), np.array(scores))
            keep_inds = list(np.array(values['indexes'])[keep_inds])
            all_keep_inds.extend(keep_inds)

        preds_filtered = {
            'predictions': pred_od,
            'hoi_prediction': list(np.array(preds['hoi_prediction'])[all_keep_inds])}
        return preds_filtered

    def pairwise_nms(self, subs, objs, scores):
        sx1, sy1, sx2, sy2 = subs[:, 0], subs[:, 1], subs[:, 2], subs[:, 3]
        ox1, oy1, ox2, oy2 = objs[:, 0], objs[:, 1], objs[:, 2], objs[:, 3]

        sub_areas = (sx2 - sx1 + 1) * (sy2 - sy1 + 1)
        obj_areas = (ox2 - ox1 + 1) * (oy2 - oy1 + 1)

        order = scores.argsort()[::-1]

        keep_inds = []
        while order.size > 0:
            i = order[0]
            keep_inds.append(i)

            sxx1 = np.maximum(sx1[i], sx1[order[1:]])
            syy1 = np.maximum(sy1[i], sy1[order[1:]])
            sxx2 = np.minimum(sx2[i], sx2[order[1:]])
            syy2 = np.minimum(sy2[i], sy2[order[1:]])

            sw = np.maximum(0.0, sxx2 - sxx1 + 1)
            sh = np.maximum(0.0, syy2 - syy1 + 1)
            sub_inter = sw * sh
            sub_union = sub_areas[i] + sub_areas[order[1:]] - sub_inter

            oxx1 = np.maximum(ox1[i], ox1[order[1:]])
            oyy1 = np.maximum(oy1[i], oy1[order[1:]])
            oxx2 = np.minimum(ox2[i], ox2[order[1:]])
            oyy2 = np.minimum(oy2[i], oy2[order[1:]])

            ow = np.maximum(0.0, oxx2 - oxx1 + 1)
            oh = np.maximum(0.0, oyy2 - oyy1 + 1)
            obj_inter = ow * oh
            obj_union = obj_areas[i] + obj_areas[order[1:]] - obj_inter

            ovr = np.power(sub_inter/sub_union, self.nms_alpha) * np.power(obj_inter / obj_union, self.nms_beta)
            inds = np.where(ovr <= self.nms_thresh)[0]

            order = order[inds + 1]
        return keep_inds


    def voc_ap(self, rec, prec):
        ap = 0.
        for t in np.arange(0., 1.1, 0.1):
            if np.sum(rec >= t) == 0:
                p = 0
            else:
                p = np.max(prec[rec >= t])
            ap = ap + p / 11.
        return ap
#############################################  PNMS ###########################################################

